// 本代码由Deepseek生成
#include "vector/fashion_mnist_load.h"

#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <cstdint>
#include <algorithm>
#include "vector/one_hot.h"
#include "swap_endian.h"

FashionMnistLoad::FashionMnistLoad(const std::string& resource_dir) : resource_dir_(resource_dir)
{
    
}

int64_t FashionMnistLoad::load()
{
    try {
        train_images_ = load_images(resource_dir_ + "train-images-idx3-ubyte");
        train_labels_ = load_labels(resource_dir_ + "train-labels-idx1-ubyte");
        test_images_ = load_images(resource_dir_ + "t10k-images-idx3-ubyte");
        test_labels_ = load_labels(resource_dir_ + "t10k-labels-idx1-ubyte");

        normalized_train_images_ = normalize(train_images_);
        one_hot_train_labels_ = OneHot::encode(train_labels_);
        normalized_test_images_ = normalize(test_images_);
        one_hot_test_labels_ = OneHot::encode(test_labels_);
    } catch (const std::exception& e) {
        std::cerr << "错误: " << e.what() << "\n";
        return 1;
    }

    return 0;
}

const std::vector<std::vector<uint8_t>>& FashionMnistLoad::get_train_images()
{
    return train_images_;
}
const std::vector<std::vector<double>>& FashionMnistLoad::get_normalized_train_images()
{
    return normalized_train_images_;
}
const std::vector<uint8_t>& FashionMnistLoad::get_train_labels()
{
    return train_labels_;
}
const std::vector<std::vector<uint8_t>>& FashionMnistLoad::get_one_hot_train_labels()
{
    return one_hot_train_labels_;
}
const std::vector<std::vector<uint8_t>>& FashionMnistLoad::get_test_images()
{
    return test_images_;
}
const std::vector<std::vector<double>>& FashionMnistLoad::get_normalized_test_images()
{
    return normalized_test_images_;
}
const std::vector<uint8_t>& FashionMnistLoad::get_test_labels()
{
    return test_labels_;
}
const std::vector<std::vector<uint8_t>>& FashionMnistLoad::get_one_hot_test_labels()
{
    return one_hot_test_labels_;
}

// 加载图像数据
std::vector<std::vector<uint8_t>> FashionMnistLoad::load_images(const std::string& path, int max_items) {
    std::ifstream file(path, std::ios::binary);
    if (!file) {
        throw std::runtime_error("无法打开文件: " + path);
    }

    uint32_t magic, num_images, rows, cols;
    file.read(reinterpret_cast<char*>(&magic), 4);
    file.read(reinterpret_cast<char*>(&num_images), 4);
    file.read(reinterpret_cast<char*>(&rows), 4);
    file.read(reinterpret_cast<char*>(&cols), 4);

    // 转换为小端序
    magic = SwapEndian::swap(magic);
    num_images = SwapEndian::swap(num_images);
    rows = SwapEndian::swap(rows);
    cols = SwapEndian::swap(cols);

    if (magic != 0x803) {
        throw std::runtime_error("无效的图像文件格式");
    }

    if (max_items > 0 && num_images > static_cast<uint32_t>(max_items)) {
        num_images = max_items;
    }

    std::vector<std::vector<uint8_t>> images(num_images, std::vector<uint8_t>(rows * cols));
    for (uint32_t i = 0; i < num_images; ++i) {
        file.read(reinterpret_cast<char*>(images[i].data()), rows * cols);
    }

    return images;
}

// 加载标签数据
std::vector<uint8_t> FashionMnistLoad::load_labels(const std::string& path, int max_items) {
    std::ifstream file(path, std::ios::binary);
    if (!file) {
        throw std::runtime_error("无法打开文件: " + path);
    }

    uint32_t magic, num_labels;
    file.read(reinterpret_cast<char*>(&magic), 4);
    file.read(reinterpret_cast<char*>(&num_labels), 4);

    // 转换为小端序
    magic = SwapEndian::swap(magic);
    num_labels = SwapEndian::swap(num_labels);

    if (magic != 0x801) {
        throw std::runtime_error("无效的标签文件格式");
    }

    if (max_items > 0 && num_labels > static_cast<uint32_t>(max_items)) {
        num_labels = max_items;
    }

    std::vector<uint8_t> labels(num_labels);
    file.read(reinterpret_cast<char*>(labels.data()), num_labels);

    return labels;
}

std::vector<std::vector<double>> FashionMnistLoad::normalize(const std::vector<std::vector<uint8_t>>& input) {
    std::vector<std::vector<double>> normalized;
    normalized.reserve(input.size()); // 预分配空间以提高性能

    for (const auto& row : input) {
        std::vector<double> normalized_row;
        normalized_row.reserve(row.size()); // 预分配行空间

        for (uint8_t pixel : row) {
            normalized_row.push_back(static_cast<double>(pixel) / 255.0);
        }

        normalized.push_back(std::move(normalized_row)); // 避免拷贝
    }

    return normalized;
}